import tensorflow as tf
import numpy as np
import gzip
import struct

train_images_file = "MNIST_data/train-images-idx3-ubyte.gz"
train_labels_file = "MNIST_data/train-labels-idx1-ubyte.gz"
t10k_images_file = "MNIST_data/t10k-images-idx3-ubyte.gz"
t10k_labels_file = "MNIST_data/t10k-labels-idx1-ubyte.gz"

def read_labels(filename):
	labels = []
	
	with gzip.open(filename) as bytestream:
		index = 0
		buf = bytestream.read()
		bytestream.close()        # 根据MINIST文件的描述，文件开始是用于校验的数字，`integer`格式，占用4个字节，位于0-4位置
		# 第二个描述文件的内容数量，`integer`格式，占用4个字节，位置4-8位置
		magic, numberOfLabels = struct.unpack_from('>II', buf, index)
		print(magic)
		print(numberOfLabels)        # index += struct.calcsize('>II') #这里的结果是 +=8，为了直观，就直接填写8
		# 因为magic, numberOfLabels 占据前面8个字节，所以把下标移动到第 8 位，开始读取数字标签的内容
		index = 8
		while len(labels) < 4000:            # 根据MINIST文件的描述，labels的数字是`unsigned byte`格式，占用一个字节，所以这里填写`B`
			num = int(struct.unpack_from('B', buf, index)[0])
			tmp =[0,0,0,0,0,0,0,0,0,0]
			tmp[num] = 1
			labels.append(tmp)            # index += struct.calcsize('B')
			# 移动到下一个光标
			index += 1
	return np.array(labels)
	
def read_images(filename, labels):
	images = []	# 把文件解压成字节流
	with gzip.open(filename) as bytestream:
		index = 0
		buf = bytestream.read()
		bytestream.close()        # 根据MINIST文件的描述，文件开始是用于校验的数字，`integer`格式，占用4个字节，位于0-4位置
		# 第二个描述文件的内容数量，`integer`格式，占用4个字节，位置4-8位置
		magic, numberOfImages, rows, columns = struct.unpack_from('>IIII', buf, index)
		print("read_images magic ",magic)
		print("read_images numberOfImages", numberOfImages)
		print("read_images rows", rows)
		print("read_images columns", columns)        # index += struct.calcsize('>IIII') #这里的结果是 +=16，为了直观，就直接填写16
		# 因为magic, numberOfImages, rows, columns 占据前面16个字节，所以把下标移动到第 16 位，开始读取数字标签的内容
		index = 16

		for i in range(4000):		# 打印对应的数字标签
			image = []
			
			for x in range(rows):
				row = []
				
				for y in range(columns):
					num = int(struct.unpack_from('B', buf, index)[0])
					image.append(float(num/255.0))
					index += 1
					
			images.append(image)

	return np.array(images)# 解析labels的内容，train_labels包含了60000个数字标签，返回60000个数字标签的数组

def unpack_custom_image(filename):
	
	r = []
	text = open(filename)
	try:
		file_context = text.read()

		for item in file_context.split(','):
			r.append(int(item))

	finally:
		text.close()

	return np.array(r)

train_labels = read_labels(train_labels_file)
train_images = read_images(train_images_file, train_labels)

test_labels = read_labels(t10k_labels_file)
test_images = read_images(t10k_images_file, test_labels)

x_ = tf.placeholder("float", [784])
y_ = tf.placeholder("float")
W = tf.Variable(tf.zeros([784, 10]), name="W")
b = tf.Variable(tf.zeros([10]), name="b")
y = tf.nn.softmax(tf.matmul(tf.reshape(x_, [1,784]), W) + b)
y_reshape = tf.reshape(y, [10])

cross_entropy = -tf.reduce_sum(y_ * tf.log(y_reshape))
train_step = tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy)
correct_prediction = tf.equal(tf.argmax(y_reshape, 0), tf.argmax(y_, 0))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))


with tf.Session() as sess:
	
	sess.run(tf.global_variables_initializer())
	
	#train
	for i in range(1000):
		print("train i: ", i)
		batch_xs = train_images[i % 4000]
		batch_ys = train_labels[i % 4000]
		
		(_, W_r, b_r) = sess.run([train_step, W, b], feed_dict={x_: batch_xs, y_: batch_ys})

		#(y_r,_,_) = sess.run([y_reshape, W, b], feed_dict={x_: batch_xs, y_: batch_ys})
		#print("y_reshape:", y_r, " batch_ys: ", batch_ys)
		#print("argmax1: ", sess.run(tf.argmax(y_r, 0)), "argmax2: ", sess.run(tf.argmax(batch_ys, 0)))

		#print("W_r", sess.run(tf.reduce_sum(W_r)))
		print("b_r", b_r)
		#print()

	#test
	# for i in range(20):
		# print("test i: ", i)
		# batch_xs = test_images[i % 100]
		# batch_ys = test_labels[i % 100]
		
		# (_, W_r, b_r) = sess.run([train_step, W, b], feed_dict={x_: batch_xs, y_: batch_ys})

		# (y_r,_,_) = sess.run([y_reshape, W, b], feed_dict={x_: batch_xs, y_: batch_ys})
		# print("y_reshape:", y_r, " batch_ys: ", batch_ys)
		# print("argmax1: ", sess.run(tf.argmax(y_r, 0)), "argmax2: ", sess.run(tf.argmax(batch_ys, 0)))
		# accuracy_r,_,_ = sess.run([accuracy, W, b], feed_dict={x_: batch_xs, y_: batch_ys})

		# print("accuracy_r", accuracy_r)
		# print()

	#custom test
	r = unpack_custom_image("test_image.txt")
	print(r.shape)
	correct = np.array([1,0,0,0,0,0,0,0,0,0])
	(_, W_r, b_r) = sess.run([train_step, W, b], feed_dict={x_: r, y_: correct})
	(y_r,_,_) = sess.run([y_reshape, W, b], feed_dict={x_: r, y_: correct})
	print("y_reshape:", y_r, " batch_ys: ", correct)
	print("argmax1: ", sess.run(tf.argmax(y_r, 0)), "argmax2: ", sess.run(tf.argmax(correct, 0)))
	accuracy_r,_,_ = sess.run([accuracy, W, b], feed_dict={x_: r, y_: correct})

	print("accuracy_r", accuracy_r)